# 导入PyTorch核心库：用于构建神经网络、张量计算和自动求导
import torch
# 从torchvision导入数据集、预训练模型和数据预处理工具
from torchvision import datasets, transforms
# 从torch.utils.data导入DataLoader：用于批量加载数据
from torch.utils.data import DataLoader
# 导入PyTorch的函数式API：提供激活函数等无状态操作
import torch.nn.functional as F
# 导入PyTorch优化器库：提供SGD等优化算法
import torch.optim as optim
# 导入matplotlib绘图库：用于绘制训练损失和准确率曲线
import matplotlib.pyplot as plt

# 设置中文显示：解决matplotlib绘图时中文乱码问题
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False

# 定义超参数
batch_size = 64        # 每批次处理的样本数量
learning_rate = 0.01   # 学习率
momentum = 0.5         # 动量参数
epochs = 10            # 训练轮次

# 定义数据预处理流水线
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转为PyTorch张量并归一化到[0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # 使用MNIST数据集的均值和标准差进行标准化
])

# 加载MNIST训练集
train_dataset = datasets.MNIST(
    root='../dataset/mnist',
    train=True,
    transform=transform,
    download=True
)

# 构建训练集DataLoader
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True  # 训练时打乱数据
)

# 加载MNIST测试集
test_dataset = datasets.MNIST(
    root='../dataset/mnist',
    train=False,
    transform=transform,
    download=True
)

# 构建测试集DataLoader
test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=batch_size,
    shuffle=False  # 测试时不需要打乱数据
)


class Net(torch.nn.Module):
    """定义卷积神经网络模型"""
    def __init__(self):
        super(Net, self).__init__()
        # 第一个卷积层：1个输入通道，10个输出通道，5x5卷积核
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)
        # 第二个卷积层：10个输入通道，20个输出通道，5x5卷积核
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
        # 最大池化层：2x2的池化窗口
        self.pooling = torch.nn.MaxPool2d(2)
        # 全连接层：320个输入特征，10个输出（对应10个数字）
        self.fc1 = torch.nn.Linear(320, 10)

    def forward(self, x):
        """前向传播函数"""
        batch_size = x.size(0)  # 获取批次大小

        # 第一层卷积 -> 池化 -> ReLU激活
        x = F.relu(self.pooling(self.conv1(x)))
        # 第二层卷积 -> 池化 -> ReLU激活
        x = F.relu(self.pooling(self.conv2(x)))
        # 展平特征图，为全连接层做准备
        x = x.view(batch_size, -1)  # 结果为(batch_size, 320)
        # 全连接层 -> ReLU激活
        x = self.fc1(x)

        return x


# 选择合适的计算设备
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(f"使用计算设备: {device}")

# 初始化模型并移动到指定设备
model = Net().to(device)

# 定义损失函数（交叉熵损失，适用于分类问题）
criterion = torch.nn.CrossEntropyLoss()
# 定义优化器（带动量的SGD）
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)


def train(epoch):
    """训练函数：执行一轮训练并返回平均损失"""
    model.train()  # 设置模型为训练模式
    running_loss = 0.0
    total_batches = len(train_loader)

    for batch_idx, (inputs, target) in enumerate(train_loader):
        # 将数据移动到指定设备
        inputs, target = inputs.to(device), target.to(device)

        # 梯度清零
        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, target)

        # 反向传播和参数更新
        loss.backward()
        optimizer.step()

        # 累加损失
        running_loss += loss.item()

        # 每100个批次打印一次中间结果
        if batch_idx % 100 == 0:
            3

    # 计算并返回本轮平均损失
    avg_loss = running_loss / total_batches
    return avg_loss


def test():
    """测试函数：在测试集上评估模型并返回准确率"""
    model.eval()  # 设置模型为评估模式
    correct = 0
    total = 0

    # 测试时不需要计算梯度
    with torch.no_grad():
        for inputs, target in test_loader:
            # 将数据移动到指定设备
            inputs, target = inputs.to(device), target.to(device)

            # 前向传播获取预测结果
            outputs = model(inputs)
            # 获取预测概率最大的类别
            _, predicted = torch.max(outputs.data, 1)

            # 统计总样本数和正确预测数
            total += target.size(0)
            correct += (predicted == target).sum().item()

    # 计算准确率
    accuracy = 100 * correct / total
    print(f'测试准确率: {accuracy:.2f}% ({correct}/{total})\n')
    return accuracy


if __name__ == '__main__':
    # 初始化记录列表
    epoch_list = []
    loss_list = []
    accuracy_list = []

    # 训练循环
    for epoch in range(epochs):
        # 训练模型
        avg_loss = train(epoch)
        # 测试模型
        accuracy = test()

        # 记录数据
        epoch_list.append(epoch + 1)
        loss_list.append(avg_loss)
        accuracy_list.append(accuracy)

    # 绘制训练损失和测试准确率曲线
    fig, ax1 = plt.subplots(figsize=(10, 6))

    # 绘制训练损失曲线
    ax1.plot(epoch_list, loss_list, 'b-', linewidth=2, label='训练损失')
    ax1.set_xlabel('训练轮次')
    ax1.set_ylabel('训练损失', color='b')
    ax1.tick_params(axis='y', labelcolor='b')

    # 创建第二个y轴用于绘制准确率
    ax2 = ax1.twinx()
    ax2.plot(epoch_list, accuracy_list, 'r-', linewidth=2, label='测试准确率')
    ax2.set_ylabel('测试准确率 (%)', color='r')
    ax2.tick_params(axis='y', labelcolor='r')
    ax2.set_ylim(0, 100)  # 准确率范围限制在0-100%

    # 添加图例和标题
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='best')
    plt.title('MNIST分类模型训练损失与测试准确率')

    # 显示图像
    plt.tight_layout()
    plt.show()
